from typing import Optional, List

import numpy as np
import pandas as pd
import pywt

from model.FaultDetectionConfig import FaultDetectionConfig


class WaveletProcessor:
    """小波变换处理器"""

    def __init__(self, config: FaultDetectionConfig):
        self.config = config

    def denoising(self, series: pd.Series, method: str = 'soft') -> pd.Series:
        """
        小波去噪

        Args:
            series: 原始信号
            method: 'soft' 或 'hard' 阈值方法

        Returns:
            去噪后的信号
        """
        values = series.values
        wavelet = self.config.wavelet
        level = pywt.dwt_max_level(len(values), wavelet)

        # 小波分解
        coeffs = pywt.wavedec(values, wavelet, level=level)

        # 计算阈值（使用MAD方法）
        sigma = np.median(np.abs(coeffs[-1])) / 0.6745
        threshold = sigma * np.sqrt(2 * np.log(len(values)))

        # 阈值处理
        coeffs_thresh = list(coeffs)
        for i in range(1, len(coeffs)):
            if method == 'soft':
                coeffs_thresh[i] = pywt.threshold(coeffs[i], threshold, 'soft')
            else:
                coeffs_thresh[i] = pywt.threshold(coeffs[i], threshold, 'hard')

        # 重构信号
        denoised = pywt.waverec(coeffs_thresh, wavelet)[:len(values)]

        return pd.Series(denoised, index=series.index)